package com.zw.chat.service.impl;

import com.zw.chat.common.context.MessageContext;
import com.zw.chat.common.sse.SseSession;
import com.zw.chat.domain.ChatMessage;
import com.zw.chat.domain.ChatSession;
import com.zw.chat.repository.ChatMessageRepository;
import com.zw.chat.repository.ChatSessionRepository;
import com.zw.chat.service.IChatMessageService;
import lombok.SneakyThrows;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.domain.PageRequest;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.Date;
import java.util.LinkedList;
import java.util.List;
import java.util.Objects;

/**
 * @author by: zyl
 * @ClassName: ChatMessageServiceImpl
 * @Description: 聊天消息服务实现类
 * @Date: 2024/12/5 11:45
 */
@Service
public class ChatMessageServiceImpl implements IChatMessageService {

    @Autowired
    private ChatMessageRepository chatMessageRepository;
    @Autowired
    private ChatSessionRepository chatSessionRepository;
    @Autowired
    private ChatClient chatClient;

    /* 最大长度限制 */
    private final Integer MAX_LEN = 20;

    @Override
    public List<ChatMessage> queryChatMessage(String sessionId, int page, int size) {
        PageRequest pageRequest = PageRequest.of(page, size);
        return chatMessageRepository.queryChatMessageBySessionId(sessionId, pageRequest);
    }

    @Override
    public ChatMessage saveOrUpdateChatMessage(ChatMessage message) {
        return chatMessageRepository.save(message);
    }

    @Override
    public void deleteChatMessageById(String id) {
        chatMessageRepository.deleteById(id);
    }

    @Override
    public ChatMessage findChatMessageById(String id) {
        return chatMessageRepository.findById(id).orElse(null);
    }

    @SneakyThrows
    @Override
    public void sendMessage(String sessionId) {
        // 1.获取sse
        SseEmitter sseEmitter = SseSession.getEmitter(sessionId);
        if (Objects.isNull(sseEmitter)) {
            return;
        }
        // 2.获取会话session
        ChatSession chatSession = chatSessionRepository.findById(sessionId).orElse(null);
        if (Objects.isNull(chatSession)) {
            return;
        }

        // 3.构建请求消息
        LinkedList<Message> historyMessage = MessageContext.getMessage(sessionId);
        Message userMessage = historyMessage.getLast();
        ChatMessage requestMessage = new ChatMessage();
        requestMessage.setRole("user");
        requestMessage.setContent(userMessage.getContent());
        requestMessage.setSession(chatSession);
        requestMessage.setCreateTime(new Date());
        requestMessage.setUpdateTime(new Date());

        StringBuilder resContent = new StringBuilder();
        // 4.询问ChatGPT  --  联系上下文对话
        ChatClient.StreamResponseSpec stream = chatClient.prompt(new Prompt(historyMessage)).stream();

        stream.content().subscribe(item -> {
            // 5.发送响应内容
            try {
                resContent.append(item);
                sseEmitter.send(item);
            } catch (IOException e) {
                sseEmitter.complete();
                SseSession.rmEmitter(sessionId);
            }
        }, e -> {
            sseEmitter.completeWithError(e);
            SseSession.rmEmitter(sessionId);
        }, () -> {
            sseEmitter.complete();
            SseSession.rmEmitter(sessionId);
            // 6.保存消息
            chatMessageRepository.save(requestMessage);

            AssistantMessage assistantMessage = new AssistantMessage(resContent.toString());

            // 联系上下文长度限制  20条
            while (historyMessage.size() >= MAX_LEN) {
                historyMessage.removeFirst();
            }
            historyMessage.add(assistantMessage);

            ChatMessage responseMessage = new ChatMessage();
            responseMessage.setRole("assistant");
            responseMessage.setSession(chatSession);
            responseMessage.setContent(resContent.toString());
            responseMessage.setCreateTime(new Date());
            responseMessage.setUpdateTime(new Date());
            chatMessageRepository.save(responseMessage);
            // 7.更新session
            chatSession.getStatistic().plusChar(resContent.toString().length());
            chatSession.getStatistic().plusToken(resContent.toString().length());
            chatSessionRepository.save(chatSession);
        });
    }
}
